package com.math.aistudykbr.utils;

import com.alibaba.fastjson.JSON;
import com.math.aistudykbr.pojo.entity.MindMapNode;
import lombok.extern.slf4j.Slf4j;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.geom.RoundRectangle2D;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
public class MindMapJsonToPng {

    private static final int MIN_NODE_WIDTH = 100;  // 最小节点宽度
    private static final int NODE_PADDING = 20;    // 节点内边距
    private static final int MIN_NODE_HEIGHT = 40; // 最小节点高度
    private static final int H_GAP = 100;  // 增大水平间距，避免同级节点重叠
    private static final int V_GAP = 60;  // 增大垂直间距，使图更清晰
    private static final int SIBLING_MIN_DISTANCE = 30; // 增大同级节点之间的最小距离
    private static final int PADDING = 50;  // 画布边距
    private static final int MAX_TEXT_WIDTH = 180; // 文字最大宽度
    private static final int FONT_SIZE = 14;
    private static final int MAX_VISIBLE_DEPTH = 4; // 最大可见深度
    private static final double MIN_SCALE = 0.6;   // 最小缩放比例
    private static final int MAX_IMAGE_WIDTH = 3000; // 最大图像宽度
    private static final int MAX_IMAGE_HEIGHT = 2000; // 最大图像高度
    private static final String FONT_FAMILY = "微软雅黑";
    private static final Color[] COLORS = {
            new Color(59, 130, 246), new Color(34, 197, 94),
            new Color(239, 68, 68), new Color(245, 158, 11),
            new Color(139, 92, 246), new Color(236, 72, 153)
    };
    
    // 节点尺寸缓存，避免重复计算
    private static Map<MindMapNode, Dimension> nodeSizeCache = new HashMap<>();
    // 子树宽度缓存
    private static Map<MindMapNode, Integer> subtreeWidthCache = new HashMap<>();
    // 节点位置缓存
    private static Map<MindMapNode, Point> nodePositionCache = new HashMap<>();
    
    /**
     * 计算节点尺寸，考虑文本长度和换行
     */
    private static Dimension calculateNodeSize(Graphics2D g2d, MindMapNode node) {
        if (nodeSizeCache.containsKey(node)) {
            return nodeSizeCache.get(node);
        }
        
        g2d.setFont(new Font(FONT_FAMILY, Font.PLAIN, FONT_SIZE));
        FontMetrics metrics = g2d.getFontMetrics();
        
        String text = node.getContentTxt();
        if (text == null) text = "";
        
        // 如果文本太长，需要换行处理
        List<String> lines = wordWrap(g2d, text, MAX_TEXT_WIDTH);
        
        // 计算文本区域大小
        int textWidth = 0;
        for (String line : lines) {
            textWidth = Math.max(textWidth, metrics.stringWidth(line));
        }
        
        int textHeight = metrics.getHeight() * lines.size();
        
        // 计算节点尺寸，确保至少有最小宽高
        int nodeWidth = Math.max(MIN_NODE_WIDTH, textWidth + NODE_PADDING * 2);
        int nodeHeight = Math.max(MIN_NODE_HEIGHT, textHeight + NODE_PADDING * 2);
        
        Dimension size = new Dimension(nodeWidth, nodeHeight);
        nodeSizeCache.put(node, size);
        return size;
    }

    /**
     * 文本换行处理
     */
    private static List<String> wordWrap(Graphics2D g2d, String text, int maxWidth) {
        FontMetrics metrics = g2d.getFontMetrics();
        List<String> lines = new ArrayList<>();
        
        if (text == null || text.isEmpty()) {
            lines.add("");
            return lines;
        }
        
        if (metrics.stringWidth(text) <= maxWidth) {
            lines.add(text);
            return lines;
        }
        
        StringBuilder currentLine = new StringBuilder();
        char[] chars = text.toCharArray();
        
        for (int i = 0; i < chars.length; i++) {
            char c = chars[i];
            if (metrics.stringWidth(currentLine.toString() + c) <= maxWidth) {
                currentLine.append(c);
            } else {
                // 尝试在空白字符处断行
                int lastSpace = currentLine.lastIndexOf(" ");
                if (lastSpace > 0 && currentLine.length() > maxWidth / 2) {
                    // 在空格处断行
                    String line = currentLine.substring(0, lastSpace);
                    lines.add(line);
                    currentLine = new StringBuilder(currentLine.substring(lastSpace + 1)).append(c);
                } else {
                    // 直接断行
                    lines.add(currentLine.toString());
                    currentLine = new StringBuilder().append(c);
                }
            }
        }
        
        // 添加最后一行
        if (currentLine.length() > 0) {
            lines.add(currentLine.toString());
        }
        
        return lines;
    }

    /**
     * 计算子树的总宽度，考虑所有子节点的子树宽度
     */
    private static int calculateSubtreeWidth(Graphics2D g2d, MindMapNode node, int level, int maxDepth) {
        // 使用缓存
        if (subtreeWidthCache.containsKey(node)) {
            return subtreeWidthCache.get(node);
        }
        
        Dimension nodeSize = calculateNodeSize(g2d, node);
        int nodeWidth = nodeSize.width;
        
        // 如果已达到最大深度或没有子节点，返回节点自身宽度
        if (level >= maxDepth || node.getSonNode() == null || node.getSonNode().isEmpty()) {
            subtreeWidthCache.put(node, nodeWidth);
            return nodeWidth;
        }
        
        List<MindMapNode> children = node.getSonNode();
        int totalChildrenSubtreeWidth = 0;
        
        // 计算所有子节点的子树宽度总和
        for (MindMapNode child : children) {
            totalChildrenSubtreeWidth += calculateSubtreeWidth(g2d, child, level + 1, maxDepth);
        }
        
        // 添加子节点之间的间距
        if (children.size() > 1) {
            totalChildrenSubtreeWidth += H_GAP * (children.size() - 1);
        }
        
        // 总宽度取节点自身宽度和子树总宽度的较大值
        int subtreeWidth = Math.max(nodeWidth, totalChildrenSubtreeWidth);
        subtreeWidthCache.put(node, subtreeWidth);
        return subtreeWidth;
    }
    
    /**
     * 确定节点布局，使用树布局算法避免重叠
     */
    private static void calculateInitialLayout(Graphics2D g2d, MindMapNode rootNode, int maxDepth) {
        // 清理缓存
        nodeSizeCache.clear();
        subtreeWidthCache.clear();
        nodePositionCache.clear();
        
        // 预先计算所有子树的宽度
        calculateSubtreeWidth(g2d, rootNode, 0, maxDepth);
        
        // 计算根节点大小
        Dimension rootSize = calculateNodeSize(g2d, rootNode);
        int rootWidth = rootSize.width;
        int rootHeight = rootSize.height;
        
        // 设置根节点位置
        int rootX = 0;
        int rootY = 0;
        nodePositionCache.put(rootNode, new Point(rootX, rootY));
        
        // 递归布局子节点
        layoutChildren(g2d, rootNode, rootX, rootY, rootWidth, rootHeight, 0, maxDepth);
    }
    
    /**
     * 递归布局子节点，考虑子树宽度
     */
    private static void layoutChildren(Graphics2D g2d, MindMapNode node, 
                                     int x, int y, int width, int height, 
                                     int level, int maxDepth) {
        // 如果已达到最大深度或没有子节点，直接返回
        if (level >= maxDepth || node.getSonNode() == null || node.getSonNode().isEmpty()) {
            return;
        }
        
        List<MindMapNode> children = node.getSonNode();
        int childCount = children.size();
        
        // 获取每个子节点的子树宽度和自身宽度
        List<Integer> subtreeWidths = new ArrayList<>();
        List<Dimension> childSizes = new ArrayList<>();
        int totalSubtreesWidth = 0;
        
        for (MindMapNode child : children) {
            int subtreeWidth = subtreeWidthCache.get(child);
            Dimension childSize = calculateNodeSize(g2d, child);
            
            // 子树宽度和节点宽度取较大值
            int effectiveWidth = Math.max(subtreeWidth, childSize.width);
            subtreeWidths.add(effectiveWidth);
            childSizes.add(childSize);
            
            totalSubtreesWidth += effectiveWidth;
        }
        
        // 添加子节点之间的间距
        totalSubtreesWidth += (childCount - 1) * H_GAP;
        
        // 计算子节点的起始X坐标，居中对齐
        int startX;
        if (totalSubtreesWidth <= width) {
            // 如果子树总宽度小于节点宽度，则居中排列
            startX = x + (width - totalSubtreesWidth) / 2;
        } else {
            // 如果子树总宽度大于节点宽度，则居中于父节点
            startX = x - (totalSubtreesWidth - width) / 2;
        }
        
        if (startX < 0) startX = 0;
        
        // 计算子节点的Y坐标
        int childY = y + height + V_GAP;
        
        // 为每个子节点分配位置
        for (int i = 0; i < childCount; i++) {
            MindMapNode child = children.get(i);
            Dimension childSize = childSizes.get(i);
            int subtreeWidth = subtreeWidths.get(i);
            
            // 子节点的X坐标居中于其子树宽度
            int childX = startX + (subtreeWidth - childSize.width) / 2;
            
            // 将子节点位置信息存入缓存
            nodePositionCache.put(child, new Point(childX, childY));
            
            // 递归布局子节点的子节点
            layoutChildren(g2d, child, childX, childY, childSize.width, childSize.height, level + 1, maxDepth);
            
            // 更新下一个子节点的X坐标
            startX += subtreeWidth + H_GAP;
        }
        
        // 调整同级节点避免重叠
        resolveOverlaps(children, level);
    }
    
    /**
     * 解决同级节点之间的重叠问题
     */
    private static void resolveOverlaps(List<MindMapNode> siblings, int level) {
        if (siblings.size() <= 1) return;
        
        boolean hasOverlap;
        
        do {
            hasOverlap = false;
            
            // 检查每对相邻节点
            for (int i = 1; i < siblings.size(); i++) {
                MindMapNode left = siblings.get(i-1);
                MindMapNode right = siblings.get(i);
                
                Point leftPos = nodePositionCache.get(left);
                Point rightPos = nodePositionCache.get(right);
                
                // 获取左节点的子树宽度
                int leftSubtreeWidth = subtreeWidthCache.get(left);
                Dimension leftSize = nodeSizeCache.get(left);
                
                // 取左节点子树宽度和节点宽度的最大值
                int leftEffectiveWidth = Math.max(leftSubtreeWidth, leftSize.width);
                
                // 计算左节点的右边缘
                int leftRight = leftPos.x + leftEffectiveWidth;
                
                // 如果有重叠，向右移动当前节点及其后面的所有节点
                int minRequiredSpace = SIBLING_MIN_DISTANCE + H_GAP / 2; // 增大最小间距
                if (rightPos.x < leftPos.x + leftSize.width + minRequiredSpace) {
                    int shift = leftPos.x + leftSize.width + minRequiredSpace - rightPos.x;
                    
                    // 从当前节点开始，移动所有后续节点
                    for (int j = i; j < siblings.size(); j++) {
                        MindMapNode nodeToShift = siblings.get(j);
                        Point pos = nodePositionCache.get(nodeToShift);
                        nodePositionCache.put(nodeToShift, new Point(pos.x + shift, pos.y));
                        
                        // 同时移动其所有子节点
                        shiftSubtree(nodeToShift, shift, 0);
                    }
                    
                    hasOverlap = true;
                    break;
                }
            }
        } while (hasOverlap);
    }
    
    /**
     * 移动子树中的所有节点
     */
    private static void shiftSubtree(MindMapNode node, int shiftX, int shiftY) {
        if (node.getSonNode() == null) return;
        
        for (MindMapNode child : node.getSonNode()) {
            Point pos = nodePositionCache.get(child);
            if (pos != null) {
                nodePositionCache.put(child, new Point(pos.x + shiftX, pos.y + shiftY));
                shiftSubtree(child, shiftX, shiftY);
            }
        }
    }
    
    /**
     * 递归绘制节点及其子节点
     */
    private static void drawNode(Graphics2D g2d, MindMapNode node, int level, int maxDepth) {
        // 如果超过最大深度，则不绘制
        if (level > maxDepth) return;
        
        // 获取节点位置和尺寸
        Point nodePos = nodePositionCache.get(node);
        if (nodePos == null) return;
        
        int x = nodePos.x;
        int y = nodePos.y;
        
        Dimension nodeSize = nodeSizeCache.get(node);
        int nodeWidth = nodeSize.width;
        int nodeHeight = nodeSize.height;
        
        // 节点背景
        g2d.setColor(COLORS[level % COLORS.length]);
        g2d.fill(new RoundRectangle2D.Float(x, y, nodeWidth, nodeHeight, 12, 12));
        
        // 边框
        g2d.setColor(Color.WHITE);
        g2d.setStroke(new BasicStroke(2.0f));
        g2d.draw(new RoundRectangle2D.Float(x, y, nodeWidth, nodeHeight, 12, 12));
        
        // 绘制文本
        drawNodeText(g2d, node.getContentTxt(), x, y, nodeWidth, nodeHeight);
        
        // 绘制表示有折叠子节点的标记
        if (level == maxDepth && node.getSonNode() != null && !node.getSonNode().isEmpty()) {
            g2d.setColor(Color.WHITE);
            g2d.fillRoundRect(x + nodeWidth - 15, y + nodeHeight - 15, 10, 10, 5, 5);
            g2d.setColor(COLORS[level % COLORS.length]);
            g2d.setStroke(new BasicStroke(1.5f));
            g2d.drawLine(x + nodeWidth - 13, y + nodeHeight - 10, x + nodeWidth - 7, y + nodeHeight - 10);
        }
        
        // 递归绘制子节点
        if (node.getSonNode() != null && !node.getSonNode().isEmpty() && level < maxDepth) {
            for (MindMapNode child : node.getSonNode()) {
                // 绘制连接线
                drawConnectingLine(g2d, node, child);
                
                // 递归绘制子节点
                drawNode(g2d, child, level + 1, maxDepth);
            }
        }
    }
    
    /**
     * 绘制节点文本，支持多行
     */
    private static void drawNodeText(Graphics2D g2d, String text, int x, int y, int nodeWidth, int nodeHeight) {
        if (text == null) text = "";
        
        g2d.setColor(Color.WHITE);
        g2d.setFont(new Font(FONT_FAMILY, Font.BOLD, FONT_SIZE));
        FontMetrics metrics = g2d.getFontMetrics();
        
        // 获取换行后的文本
        List<String> lines = wordWrap(g2d, text, MAX_TEXT_WIDTH);
        
        // 绘制多行文本
        int lineHeight = metrics.getHeight();
        int textY = y + (nodeHeight - lineHeight * lines.size()) / 2 + metrics.getAscent();
        
        for (String line : lines) {
            int textX = x + (nodeWidth - metrics.stringWidth(line)) / 2;
            g2d.drawString(line, textX, textY);
            textY += lineHeight;
        }
    }
    
    /**
     * 绘制节点间的连接线
     */
    private static void drawConnectingLine(Graphics2D g2d, MindMapNode parent, MindMapNode child) {
        Point parentPos = nodePositionCache.get(parent);
        Point childPos = nodePositionCache.get(child);
        
        if (parentPos == null || childPos == null) return;
        
        Dimension parentSize = nodeSizeCache.get(parent);
        Dimension childSize = nodeSizeCache.get(child);
        
        int startX = parentPos.x + parentSize.width / 2;
        int startY = parentPos.y + parentSize.height;
        int endX = childPos.x + childSize.width / 2;
        int endY = childPos.y;
        
        g2d.setColor(new Color(180, 180, 180));
        g2d.setStroke(new BasicStroke(1.5f));
        
        // 绘制曲线连接
        int midY = startY + (endY - startY) / 2;
        int[] xPoints = {startX, startX, endX, endX};
        int[] yPoints = {startY, midY, midY, endY};
        
        g2d.drawPolyline(xPoints, yPoints, 4);
    }
    
    /**
     * 计算图像需要的尺寸
     */
    private static Rectangle calculateBounds() {
        if (nodePositionCache.isEmpty()) {
            return new Rectangle(0, 0, 800, 600);
        }
        
        int minX = Integer.MAX_VALUE;
        int minY = Integer.MAX_VALUE;
        int maxX = Integer.MIN_VALUE;
        int maxY = Integer.MIN_VALUE;
        
        // 遍历所有节点位置，找出边界
        for (Map.Entry<MindMapNode, Point> entry : nodePositionCache.entrySet()) {
            MindMapNode node = entry.getKey();
            Point pos = entry.getValue();
            Dimension size = nodeSizeCache.get(node);
            
            if (size != null) {
                minX = Math.min(minX, pos.x);
                minY = Math.min(minY, pos.y);
                maxX = Math.max(maxX, pos.x + size.width);
                maxY = Math.max(maxY, pos.y + size.height);
            }
        }
        
        // 添加边距
        minX -= PADDING;
        minY -= PADDING;
        maxX += PADDING;
        maxY += PADDING;
        
        return new Rectangle(minX, minY, maxX - minX, maxY - minY);
    }
    
    /**
     * 生成思维导图图片，自动选择深度
     */
    public static byte[] drawMindMapFromJson(String json) throws Exception {
        List<MindMapNode> roots = JSON.parseArray(json, MindMapNode.class);
        if (roots == null || roots.isEmpty()) throw new IllegalArgumentException("JSON无根节点");
        
        MindMapNode root = roots.get(0);
        
        // 计算最大深度
        int maxDepth = calculateMaxDepth(root);
        int visibleDepth = Math.min(maxDepth, MAX_VISIBLE_DEPTH);
        
        log.info("思维导图最大深度: {}, 可见深度: {}", maxDepth, visibleDepth);
        
        return drawMindMapWithDepth(root, visibleDepth);
    }
    
    /**
     * 计算树的最大深度
     */
    private static int calculateMaxDepth(MindMapNode node) {
        if (node.getSonNode() == null || node.getSonNode().isEmpty()) {
            return 0;
        }
        
        int maxChildDepth = 0;
        for (MindMapNode child : node.getSonNode()) {
            maxChildDepth = Math.max(maxChildDepth, calculateMaxDepth(child));
        }
        
        return maxChildDepth + 1;
    }
    
    /**
     * 生成指定深度的思维导图
     */
    private static byte[] drawMindMapWithDepth(MindMapNode root, int visibleDepth) throws Exception {
        // 创建临时图像用于计算布局
        BufferedImage tempImage = new BufferedImage(1, 1, BufferedImage.TYPE_INT_RGB);
        Graphics2D tempG2d = tempImage.createGraphics();
        
        // 计算初始布局
        calculateInitialLayout(tempG2d, root, visibleDepth);
        
        // 计算图像边界
        Rectangle bounds = calculateBounds();
        
        // 调整所有节点位置以适应画布
        for (MindMapNode node : nodePositionCache.keySet()) {
            Point pos = nodePositionCache.get(node);
            nodePositionCache.put(node, new Point(pos.x - bounds.x, pos.y - bounds.y));
        }
        
        // 确保尺寸合理
        int width = Math.min(MAX_IMAGE_WIDTH, Math.max(800, bounds.width));
        int height = Math.min(MAX_IMAGE_HEIGHT, Math.max(600, bounds.height));
        
        // 计算缩放比例，如果图像太大
        double scaleX = 1.0;
        double scaleY = 1.0;
        
        if (bounds.width > MAX_IMAGE_WIDTH) {
            scaleX = (double) MAX_IMAGE_WIDTH / bounds.width;
        }
        
        if (bounds.height > MAX_IMAGE_HEIGHT) {
            scaleY = (double) MAX_IMAGE_HEIGHT / bounds.height;
        }
        
        double scale = Math.max(Math.min(scaleX, scaleY), MIN_SCALE);
        
        // 如果需要缩放
        if (scale < 1.0) {
            for (MindMapNode node : nodePositionCache.keySet()) {
                Point pos = nodePositionCache.get(node);
                nodePositionCache.put(node, new Point((int)(pos.x * scale), (int)(pos.y * scale)));
                
                Dimension size = nodeSizeCache.get(node);
                nodeSizeCache.put(node, new Dimension((int)(size.width * scale), (int)(size.height * scale)));
            }
            
            width = (int)(width * scale);
            height = (int)(height * scale);
        }
        
        log.info("生成思维导图尺寸: {}x{}, 缩放比例: {}", width, height, scale);
        
        // 创建实际图像
        BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
        Graphics2D g2d = image.createGraphics();
        
        // 设置抗锯齿和高质量绘制
        g2d.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
        g2d.setRenderingHint(RenderingHints.KEY_TEXT_ANTIALIASING, RenderingHints.VALUE_TEXT_ANTIALIAS_ON);
        g2d.setRenderingHint(RenderingHints.KEY_RENDERING, RenderingHints.VALUE_RENDER_QUALITY);
        
        // 调整字体大小适应缩放
        if (scale < 1.0) {
            g2d.setFont(new Font(FONT_FAMILY, Font.PLAIN, (int)(FONT_SIZE * scale)));
        }
        
        // 填充背景色
        g2d.setColor(new Color(248, 250, 252));
        g2d.fillRect(0, 0, width, height);
        
        // 绘制思维导图
        drawNode(g2d, root, 0, visibleDepth);
        
        g2d.dispose();
        tempG2d.dispose();
        
        // 转换为字节数组
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        ImageIO.write(image, "PNG", out);
        return out.toByteArray();
    }
    
    /**
     * 将思维导图保存为文件
     */
    public static void drawMindMapToFile(String json, String outputPath) throws Exception {
        byte[] img = drawMindMapFromJson(json);
        File file = new File(outputPath);
        if (file.getParentFile() != null) file.getParentFile().mkdirs();
        java.nio.file.Files.write(file.toPath(), img);
        log.info("思维导图已保存到: {}", file.getAbsolutePath());
    }
} 